"""
氧化铝指数K线图和双均线副图生成工具
使用tqsdk获取数据，pyecharts生成图表
"""

import time
import pandas as pd
import numpy as np
from tqsdk import TqApi, TqAuth, TqKq
from pyecharts import options as opts
from pyecharts.charts import Kline, Line, Bar, Grid
from pyecharts.commons.utils import JsCode
import sys
import os

# 全局参数，控制需要显示的数据周期
DISPLAY_PERIOD = 500

# 导入当前模块中的算法函数
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from algorithm_module import calculate_boeyi_indicators,calculate_trading_signals

# 定义颜色常量
UP_COLOR = "#ec0000"
DOWN_COLOR = "#00da3c"
BORDER_COLOR = "#8A0000"
AREA_COLOR = "#f9f9f9"
BACKGROUND_COLOR = "#21202D"
TRADING_LINE_COLOR = "#ffa500"  # 橙色
SHORT_LINE_COLOR = "#00ffff"    # 青色

def fetch_alumina_data(days=DISPLAY_PERIOD):
    """
    获取氧化铝指数的K线数据
    
    Args:
        days: 获取多少天的数据
    
    Returns:
        pandas.DataFrame: K线数据
    """
    print("正在连接天勤获取氧化铝指数数据...")
    api = TqApi(account=TqKq(), auth=TqAuth("cps168", "alibaba"))
    
    try:
        # 获取氧化铝期货K线数据
        symbol = "DCE.p2601"  # 示例使用有效合约代码，需根据实际替换
        print(f"正在获取 {symbol} 的K线数据...")
        #klines = api.get_kline_serial(symbol, duration_seconds=60*60*24, data_length=days)
        klines = api.get_kline_serial(symbol, 60 * 60 * 24, data_length=days)
        print("K线数据获取成功，正在计算指标...")
        # 确保K线数据足够计算指标
        if len(klines) < 26:
            print(f"K线数据不足，当前只有 {len(klines)} 条，至少需要26条")
            return None
            
        # 计算操盘线和空头线
        # 计算操盘线和空头线
        # 计算指标并获取结果
        result = calculate_trading_signals(klines)

        # 合并K线数据和指标结果
        result = pd.merge(klines[['datetime', 'open', 'high', 'low', 'close', 'volume']], result, left_index=True, right_index=True)

        # 处理日期格式
        result['date'] = pd.to_datetime(result['datetime'], unit='s')
        result['date_str'] = result['date'].dt.strftime('%Y-%m-%d')
        
        print(f"成功获取 {len(result)} 条K线数据")
        return result
    
    except Exception as e:
        print(f"获取数据时出错: {e}")
        import traceback
        traceback.print_exc()
        return None
    
    finally:
        api.close()

def create_kline_chart(data):
    """
    创建K线图和双均线副图
    
    Args:
        data: 包含K线和均线数据的DataFrame
    
    Returns:
        pyecharts.charts.Grid: 图表对象
    """
    # 准备K线数据
    kline_data = data[['date_str', 'open', 'close', 'low', 'high']].values.tolist()
    
    # 准备均线数据
    trading_line_data = data[['date_str', 'trading_line']].values.tolist()
    short_line_data = data[['date_str', 'short_line']].values.tolist()
    
    # 准备成交量数据
    volume_data = data[['date_str', 'volume']].values.tolist()
    
    # 计算涨跌
    data['color'] = np.where(data['close'] > data['open'], 1, -1)
    
    # 创建K线图
    kline = (
        Kline()
        .add_xaxis(xaxis_data=data['date_str'].tolist())
        .add_yaxis(
            series_name="K线",
            y_axis=kline_data,
            itemstyle_opts=opts.ItemStyleOpts(
                color=UP_COLOR,
                color0=DOWN_COLOR,
                border_color=UP_COLOR,
                border_color0=DOWN_COLOR,
            ),
        )
        .set_global_opts(
            title_opts=opts.TitleOpts(
                title="氧化铝指数K线图",
                subtitle=f"数据周期: {data['date_str'].iloc[0]} 至 {data['date_str'].iloc[-1]}",
            ),
            xaxis_opts=opts.AxisOpts(
                type_="category",
                is_scale=True,
                boundary_gap=False,
                axisline_opts=opts.AxisLineOpts(is_on_zero=False),
                splitline_opts=opts.SplitLineOpts(is_show=False),
                split_number=20,
                min_="dataMin",
                max_="dataMax",
            ),
            yaxis_opts=opts.AxisOpts(
                is_scale=True,
                splitline_opts=opts.SplitLineOpts(is_show=True),
            ),
            tooltip_opts=opts.TooltipOpts(trigger="axis", axis_pointer_type="cross"),
            datazoom_opts=[
                opts.DataZoomOpts(
                    is_show=False,
                    type_="inside",
                    xaxis_index=[0, 1],
                    range_start=50,
                    range_end=100,
                ),
                opts.DataZoomOpts(
                    is_show=True,
                    xaxis_index=[0, 1],
                    type_="slider",
                    pos_top="85%",
                    range_start=50,
                    range_end=100,
                ),
            ],
            legend_opts=opts.LegendOpts(
                is_show=True, pos_bottom=10, pos_left="center"
            ),
        )
    )
    
    # 创建均线图，添加到K线图上
    line = (
        Line()
        .add_xaxis(xaxis_data=data['date_str'].tolist())
        .add_yaxis(
            series_name="操盘线",
            y_axis=[item[1] for item in trading_line_data],
            is_smooth=True,
            is_symbol_show=False,
            linestyle_opts=opts.LineStyleOpts(width=2, color="#FFA500"),
            label_opts=opts.LabelOpts(is_show=False),
        )
        .add_yaxis(
            series_name="空头线",
            y_axis=[item[1] for item in short_line_data],
            is_smooth=True,
            is_symbol_show=False,
            linestyle_opts=opts.LineStyleOpts(width=2, color="#00FFFF"),
            label_opts=opts.LabelOpts(is_show=False),
        )
        .set_global_opts(
            xaxis_opts=opts.AxisOpts(
                type_="category",
                grid_index=1,
                axislabel_opts=opts.LabelOpts(is_show=False),
            ),
            yaxis_opts=opts.AxisOpts(
                grid_index=1,
                split_number=4,
                axisline_opts=opts.AxisLineOpts(is_on_zero=False),
                axistick_opts=opts.AxisTickOpts(is_show=False),
                splitline_opts=opts.SplitLineOpts(is_show=False),
                axislabel_opts=opts.LabelOpts(is_show=True),
            ),
        )
    )
    
    # 添加金叉和死叉标注
    cross_points = []
    for i in range(1, len(data)):
        if data['trading_line'].iloc[i] > data['short_line'].iloc[i] and data['trading_line'].iloc[i-1] <= data['short_line'].iloc[i-1]:
            cross_points.append({
                'coord': [data['date_str'].iloc[i], data['trading_line'].iloc[i]],
                'symbol': 'triangle',
                'symbol_size': 10,
                'itemstyle_opts': opts.ItemStyleOpts(color='#00ff00'),
                'label_opts': opts.LabelOpts(
                    formatter='金叉',
                    position='top',
                    color='#00ff00'
                )
            })
        elif data['trading_line'].iloc[i] < data['short_line'].iloc[i] and data['trading_line'].iloc[i-1] >= data['short_line'].iloc[i-1]:
            cross_points.append({
                'coord': [data['date_str'].iloc[i], data['trading_line'].iloc[i]],
                'symbol': 'triangle',
                'symbol_size': 10,
                'itemstyle_opts': opts.ItemStyleOpts(color='#ff0000'),
                'label_opts': opts.LabelOpts(
                    formatter='死叉',
                    position='bottom',
                    color='#ff0000'
                )
            })
    
    # 添加金叉和死叉标注
    js_func = """
    function (params) {
        return {
            type: 'scatter',
            coordinateSystem: 'cartesian2d',
            symbol: params.symbol,
            symbolSize: params.symbol_size,
            itemStyle: params.itemstyle_opts,
            label: params.label_opts,
            data: [params.coord]
        };
    }
    """
    import json
    for point in cross_points:
        new_point = point.copy()
        new_point['itemstyle_opts'] = point['itemstyle_opts'].__dict__
        new_point['label_opts'] = point['label_opts'].__dict__
        line.add_js_funcs(js_func, json.dumps(new_point))
    
    # 创建成交量图
    bar = (
        Bar()
        .add_xaxis(xaxis_data=data['date_str'].tolist())
        .add_yaxis(
            series_name="成交量",
            y_axis=[item[1] for item in volume_data],
            xaxis_index=1,
            yaxis_index=1,
            label_opts=opts.LabelOpts(is_show=False),
            itemstyle_opts=opts.ItemStyleOpts(
                color=JsCode(
                    """
                    function(params) {
                        var colorList;
                        if (params.data >= 0) {
                            colorList = '#ef232a';
                        } else {
                            colorList = '#14b143';
                        }
                        return colorList;
                    }
                    """
                )
            ),
        )
        .set_global_opts(
            xaxis_opts=opts.AxisOpts(
                type_="category",
                grid_index=1,
                axislabel_opts=opts.LabelOpts(is_show=False),
            ),
            legend_opts=opts.LegendOpts(is_show=False),
        )
    )
    
    # 创建双均线副图
    line_sub = (
        Line()
        .add_xaxis(xaxis_data=data['date_str'].tolist())
        .add_yaxis(
            series_name="操盘线",
            y_axis=[item[1] for item in trading_line_data],
            xaxis_index=2,
            yaxis_index=2,
            is_smooth=True,
            is_symbol_show=False,
            linestyle_opts=opts.LineStyleOpts(width=2, color="#FFA500"),
            label_opts=opts.LabelOpts(is_show=False),
        )
        .add_yaxis(
            series_name="空头线",
            y_axis=[item[1] for item in short_line_data],
            xaxis_index=2,
            yaxis_index=2,
            is_smooth=True,
            is_symbol_show=False,
            linestyle_opts=opts.LineStyleOpts(width=2, color="#00FFFF"),
            label_opts=opts.LabelOpts(is_show=False),
        )
        .set_global_opts(
            title_opts=opts.TitleOpts(
                title="双均线指标",
                pos_left="center",
                pos_top="bottom",
            ),
            xaxis_opts=opts.AxisOpts(
                type_="category",
                grid_index=2,
                axislabel_opts=opts.LabelOpts(is_show=True),
                split_number=20,
                min_="dataMin",
                max_="dataMax",
            ),
            yaxis_opts=opts.AxisOpts(
                grid_index=2,
                split_number=4,
                axisline_opts=opts.AxisLineOpts(is_on_zero=False),
                axistick_opts=opts.AxisTickOpts(is_show=True),
                splitline_opts=opts.SplitLineOpts(is_show=True),
                axislabel_opts=opts.LabelOpts(is_show=True),
            ),
            legend_opts=opts.LegendOpts(
                is_show=True, pos_bottom=50, pos_left="center"
            ),
        )
    )
    
    # 组合图表
    grid = (
        Grid()
        .add(
            kline,
            grid_opts=opts.GridOpts(
                pos_left="10%", pos_right="8%", height="50%"
            ),
        )
        .add(
            line,
            grid_opts=opts.GridOpts(
                pos_left="10%", pos_right="8%", pos_top="55%", height="20%"
            ),
        )
        .add(
            bar,
            grid_opts=opts.GridOpts(
                pos_left="10%", pos_right="8%", pos_top="80%", height="10%"
            ),
        )
    )
    
    return grid

def generate_mock_data(days=60):
    """
    生成模拟数据用于测试
    
    Args:
        days: 生成多少天的数据
    
    Returns:
        pandas.DataFrame: 模拟的K线数据
    """
    print("生成模拟数据用于测试...")
    
    # 生成日期序列
    import datetime
    end_date = datetime.datetime.now()
    start_date = end_date - datetime.timedelta(days=days)
    date_range = pd.date_range(start=start_date, end=end_date, freq='D')
    
    # 生成模拟价格数据
    np.random.seed(42)  # 设置随机种子，确保结果可重现
    
    # 基础价格和波动范围
    base_price = 15000  # 铝的基础价格
    volatility = 500    # 价格波动范围
    
    # 生成开盘价、收盘价、最高价、最低价
    closes = base_price + np.cumsum(np.random.normal(0, 100, size=len(date_range)))
    opens = closes + np.random.normal(0, 50, size=len(date_range))
    highs = np.maximum(opens, closes) + np.random.uniform(0, 100, size=len(date_range))
    lows = np.minimum(opens, closes) - np.random.uniform(0, 100, size=len(date_range))
    volumes = np.random.randint(1000, 10000, size=len(date_range))
    
    # 计算模拟的操盘线和空头线
    # 操盘线：26日加权移动平均
    trading_line = pd.Series(closes).rolling(window=26).mean()
    # 空头线：操盘线的7日指数移动平均
    short_line = pd.Series(trading_line).ewm(span=7).mean()
    
    # 创建DataFrame
    data = pd.DataFrame({
        'date': date_range,
        'date_str': [d.strftime('%Y-%m-%d') for d in date_range],
        'open': opens,
        'high': highs,
        'low': lows,
        'close': closes,
        'volume': volumes,
        'trading_line': trading_line,
        'short_line': short_line
    })
    
    # 填充NaN值
    data = data.fillna(method='ffill').fillna(method='bfill')
    
    print(f"成功生成 {len(data)} 条模拟K线数据")
    return data

def main():
    """
    主函数
    """
    print("氧化铝指数K线图和双均线副图生成工具")
    print("=" * 50)
    
    try:
        # 尝试获取真实数据
        days = 60
        print(f"准备获取最近 {days} 天的数据...")
        data = fetch_alumina_data(days)
        
        # 如果获取真实数据失败，则使用模拟数据
        if data is None or len(data) == 0:
            print("未能获取有效数据，将使用模拟数据进行演示")
            data = generate_mock_data(days)
            
            if data is None or len(data) == 0:
                print("生成模拟数据也失败了，程序退出")
                return
        
        # 检查数据是否包含必要的列
        required_columns = ['date_str', 'open', 'close', 'low', 'high', 'trading_line', 'short_line', 'volume']
        missing_columns = [col for col in required_columns if col not in data.columns]
        if missing_columns:
            print(f"数据缺少必要的列: {missing_columns}")
            print(f"可用的列: {data.columns.tolist()}")
            return
            
        # 检查数据中是否有NaN值
        nan_columns = [col for col in required_columns if data[col].isna().any()]
        if nan_columns:
            print(f"以下列中包含NaN值: {nan_columns}")
            print("尝试填充NaN值...")
            for col in nan_columns:
                data[col] = data[col].fillna(method='ffill').fillna(method='bfill')
        
        # 创建图表
        print("正在生成图表...")
        chart = create_kline_chart(data)
        
        # 保存图表
        output_file = "alumina_chart.html"
        chart.render(output_file)
        print(f"图表已保存至: {os.path.abspath(output_file)}")
        
        # 自动打开图表
        try:
            import webbrowser
            webbrowser.open(output_file)
            print("已自动打开图表")
        except Exception as e:
            print(f"无法自动打开图表: {e}")
            print(f"请手动打开文件: {os.path.abspath(output_file)}")
    
    except Exception as e:
        print(f"程序执行过程中出错: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()